"""Flux for the Burger's equation."""

from hogs.solvers.flux.flux_functions import FluxFunction, WAFFlux
from hogs.solvers.riemann.burgers import burgers

import numpy

class Godunov(FluxFunction):
    """Godunov's flux for the Burger's equation."""

    def __init__(self):
        self.rsolver = burgers.BurgersExact()

    def compute_flux(self, grid, dt):
        for face in grid.faces[grid.nb:-grid.nb]:
            ql, qr = face.get_conserved_variables(i=0)

            # get the Godunov's solution for the input data
            q = self.rsolver.godunov_solution(ql, qr)

            face.flux[:] = 0.5 * q * q

class WAF(WAFFlux):
    """Weighted average flux for Burger's equation"""

    def __init__(self, limiter_type=0):
        WAFFlux.__init__(self, limiter_type)
        self.rsolver = burgers.BurgersExact()
        
    def _get_data_state_vector(self):
        return numpy.zeros( shape=(1, 2) )

    def _get_wave_speed_vector(self):
        return numpy.zeros(1)

    def _get_jump(self, ql, qr):
        return qr[0] - ql[0]

    def _compute_waf_states(self, ql, qr, face, U, S):
        self.rsolver.compute_waf_states(ql, qr, U, S)

    def _compute_unstable_flux(self, ql, qr):
        return 0.5 * ( burgers.burgers_flux(ql) + \
                       burgers.burgers_flux(qr) )

    def _compute_delF(self, U, i):
        return burgers.burgers_flux(U[:, 1]) - \
               burgers.burgers_flux(U[:, 0])
